#include <cstdio> 
#include "mpi.h" 
#include <cstdlib>
// #include <xpmem.h>

// #define YIZIMI_SIMPLE_ALLREDUCE
// #define YIZIMI_RING_ALLREDUCE
#define YIZIMI_XPMEM_RING
// #define PROFILE_PRINT

static double totalTime = 0.0;
static int totalBytes = 0;
static double allreduce_time = 0.00;

template <typename T> void yizimi_swap(T& a, T& b) {
    T t; t = a, a = b, b = t;
}

int MPI_Init(int *argc, char ***argv) {     

    int my_rank, err;     
    err = PMPI_Init(argc, argv);     

#ifdef PROFILE_PRINT
    PMPI_Comm_rank(MPI_COMM_WORLD, &my_rank);     
    printf("[%d]: Profile MPI_Init.\n", my_rank);    
#endif // PROFILE_PRINT 

    return err; 
}   

int MPI_Finalize() { 	

    int my_rank; 	

#ifdef PROFILE_PRINT
    PMPI_Comm_rank(MPI_COMM_WORLD, &my_rank); 	
    printf("[%d]: Profile MPI_Finalize.\n", my_rank); 	
#endif // PROFILE_PRINT

    return PMPI_Finalize(); 
}

int MPI_Allreduce(const void* sendbuf, void* recvbuf, int count, 
                  MPI_Datatype datatype, MPI_Op op, MPI_Comm comm) {

    int my_rank, comm_sz; 	
    PMPI_Comm_rank(MPI_COMM_WORLD, &my_rank); 	
    PMPI_Comm_size(MPI_COMM_WORLD, &comm_sz);

    PMPI_Barrier(comm);

    double tstart = MPI_Wtime();

    // int result = PMPI_Allreduce(sendbuf, recvbuf, count, datatype, op, comm);
    
    {
        // yizimi's allreduce
        {
        #ifdef YIZIMI_SIMPLE_ALLREDUCE
            void* getbuf;
            int res;
            if (my_rank == 0) {
                printf("[%d]: Yizimi Simple Allreduce\n", my_rank);
                if (datatype == MPI_FLOAT) {
                    getbuf = malloc(sizeof(float) * count * comm_sz);
                }
            }
            // allreduce -- 最朴素的思路：Gather + ops + Bcast
            res = PMPI_Gather(sendbuf, count, datatype, getbuf, count, datatype, 0, MPI_COMM_WORLD);
            if (res != 0) { return res; }

            if (my_rank == 0) {
                // for (int i = 0; i < comm_sz; ++i)
                //     for (int j = 0; j < count; ++j)
                //         printf("%.2lf%c", ((float*)getbuf)[i * count + j], (j == count - 1) ? '\n' : ' ');
                for (int j = 0; j < count; ++j) {
                    float tmp = 0.00;
                    for (int i = 0; i < comm_sz; ++i)
                        tmp += ((float*)getbuf)[i * count + j];
                    ((float *)recvbuf)[j] = tmp;
                }
                // for (int j = 0; j < count; ++j)
                //     printf("%.2lf%c", ((float*)recvbuf)[j], (j == count - 1) ? '\n' : ' ');
            }

            res = PMPI_Bcast(recvbuf, count, datatype, 0, MPI_COMM_WORLD);
            if (res != 0) { return res; }
        #endif // YIZIMI_SIMPLE_ALLREDUCE
        }
        {
        #ifdef YIZIMI_RING_ALLREDUCE
            // 最朴素的Ring-Allreduce
            int pre, sup;
            MPI_Status status;
            if (my_rank == 0)
                printf("[%d]: Yizimi Ring Allreduce\n", my_rank);
            pre = (my_rank - 1 + comm_sz) % comm_sz;
            sup = (my_rank + 1) % comm_sz;
            int chunk_size = count / comm_sz;
            float *tempbuf = new float[chunk_size];
            // 第一部分：Allreduce add
            for (int i = 0; i < comm_sz - 1; ++i) {
                // int send_chunk = chunk_size * ((my_rank - i + comm_sz) % comm_sz);
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] += tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
            // 第二部分：Allreduce pass
            for (int i = comm_sz - 1; i < comm_sz; ++i) {
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] = tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
            for (int i = 0; i < comm_sz - 2; ++i) {
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] = tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
        #endif // YIZIMI_RING_ALLREDUCE
        }
        {
        #ifdef YIZIMI_XPMEM_RING
            // 使用xpmem优化的Ring-Allreduce
            int pre, sup;
            MPI_Status status;
            if (my_rank == 0)
                printf("[%d]: Yizimi XPMEM Ring Allreduce\n", my_rank);
            pre = (my_rank - 1 + comm_sz) % comm_sz;
            sup = (my_rank + 1) % comm_sz;
            int chunk_size = count / comm_sz;
            float *tempbuf = new float[chunk_size];
            // 第一部分：Allreduce add
            for (int i = 0; i < comm_sz - 1; ++i) {
                // int send_chunk = chunk_size * ((my_rank - i + comm_sz) % comm_sz);
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] += tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
            // 第二部分：Allreduce pass
            for (int i = comm_sz - 1; i < comm_sz; ++i) {
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] = tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
            for (int i = 0; i < comm_sz - 2; ++i) {
                PMPI_Sendrecv((void*)((float*)(sendbuf) + chunk_size * ((my_rank - i + comm_sz) % comm_sz)), chunk_size, datatype, sup, 0,
                    (void*)(tempbuf), chunk_size, datatype, pre, 0, comm, &status);
                PMPI_Barrier(comm);
                int temp = chunk_size * ((pre - i + comm_sz) % comm_sz);
                for (int j = 0; j < chunk_size; ++j) {
                    ((float *)(sendbuf))[temp + j] = tempbuf[j],
                    ((float *)(recvbuf))[temp + j] = ((float *)(sendbuf))[temp + j];
                }
            }
        #endif // YIZIMI_XPMEM_RING
        }
    }
    allreduce_time += MPI_Wtime() - tstart;

    printf("[%d]: Profile MPI_Allreduce: tim = %lf.\n", my_rank, allreduce_time);	
    return 0;
}
